#!/bin/bash

TARGET=/run/k3os/target

source /usr/lib/os-release

grow()
{
    parted $1 resizepart $2 100%
    partprobe $1
    sleep 2
    e2fsck -f $3
    resize2fs $3
}

setup_mounts()
{
    mkdir -p $TARGET
    mount -L K3OS_STATE $TARGET

    if [ -e $TARGET/k3os/system/growpart ]; then
        read DEV NUM < $TARGET/k3os/system/growpart
        if [ ! -e "${DEV}${NUM}" ]; then
            # /dev/sda2 => /dev/sda2
            # /dev/nvme0n1p2 => /dev/nvme0n1p2
            PART=$(blkid -L K3OS_STATE)

            # /dev/sda2 => /dev/sda
            # /dev/nvme0n1p2 => /dev/nvme0n1
            DEV=$(echo "$PART" | sed -r 's/((\d+)p)?\d+$/\2/')

            # /dev/sda2 => 2
            # /dev/nvme0n1p2 => 2
            NUM=$(echo "$PART" | sed 's!.*[^0-9]!!')
        fi
        if [ -e "${PART:=${DEV}${NUM}}" ]; then
            umount $TARGET
            grow $DEV $NUM $PART || true
            mount -L K3OS_STATE $TARGET
        fi
        rm -f $TARGET/k3os/system/growpart
    fi
}

setup_kernel_squashfs()
{
    KER_SRC="/.base/k3os/system/kernel/$(uname -r)/kernel.squashfs"
    KER_DEST="$TARGET/k3os/system/kernel/$(uname -r)/kernel.squashfs"
    if [ -e $KER_SRC ] && [ ! -e $KER_DEST ]; then
        mkdir -p $(dirname ${KER_DEST})
        cp -r $KER_SRC $KER_DEST
    fi
}

setup_k3os()
{
    if [ -e $TARGET/k3os/system/k3os/current/k3os ]; then
        return 0
    fi

    K3OS_SRC=/.base/k3os/system/k3os/current/k3os
    K3OS_FILE=$TARGET/k3os/system/k3os/${VERSION_ID}/k3os

    if [ ! -e ${K3OS_SRC} ]; then
        return 0
    fi

    if [ ! -e ${K3OS_FILE} ]; then
        mkdir -p $(dirname ${K3OS_FILE}.tmp)
        cp -f ${K3OS_SRC} ${K3OS_FILE}.tmp
        mv -f ${K3OS_FILE}.tmp ${K3OS_FILE}
    fi

    ln -sf ${VERSION_ID} $TARGET/k3os/system/k3os/current
}

setup_init()
{
    if [ -e $TARGET/sbin/init ]; then
        return 0
    fi

    mkdir -p $TARGET/sbin
    ln -sf ../k3os/system/k3os/current/k3os $TARGET/sbin/init
}

setup_k3s()
{
    if [ -e $TARGET/k3os/system/k3s/current/k3s ]; then
        return 0
    fi

    for i in $(ls -drt $TARGET/k3os/system/k3s/*); do
        if [ ! -L "$i" ]; then
            LATEST=$i
            break
        fi
    done

    if [ -e "${LATEST}" ]; then
        ln -sf $(basename "${LATEST}") $TARGET/k3os/system/k3s/current
        return 0
    fi
}

takeover()
{
    if [ ! -e k3os/system/takeover ]; then
        return 0
    fi

    touch k3os/system/factory-reset

    if [ -L sbin ]; then
      rm -f sbin
    fi

    for i in *; do
        case $i in
            boot|k3os|sbin)
                ;;
            *)
                rm -rf $i
        esac
    done
    for i in sbin/*; do
        case $i in
            sbin/init|sbin/k3s|sbin/k3os)
                ;;
            *)
                rm -rf "$i"
        esac
    done

    rm -rf boot/[^g]*
    rm -f k3os/system/takeover
    rm -rf k3os/data
    sync
    
    if [ -e k3os/system/poweroff ]; then
        rm -f k3os/system/poweroff
        poweroff -f
    else
        reboot -f
    fi
}

setup_mounts
setup_k3os
setup_kernel_squashfs
setup_init
setup_k3s

cd $TARGET

takeover || true

if [ -e k3os/system/factory-reset ] || [ -e k3os/system/ephemeral ]; then
    rm -rf k3os/data k3os/system/factory-reset
fi

losetup -d /dev/loop0 || true
mount --make-rprivate /
mkdir -p .root
pivot_root . .root
K3OS_MODE=local exec /sbin/init
